import os
import sys
import json
import glob
import time
import base64
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import streamlit as st
sys.path.append(os.path.pardir)
from Data.Draw_List import DRAW_LIST
from Calculate.Calculate import Numerical, Analytic
from Plot.Plot import Plot

WORK_SPACE = os.path.abspath(os.path.join(
    os.path.dirname(__file__), os.path.pardir))
DATASETS = glob.glob(os.path.join(WORK_SPACE, 'Data', '*.csv'))
SETTINGS = glob.glob(os.path.join(WORK_SPACE, 'Data', 'Net_Setting*.json'))

def main(DATA_Cal, DATA):
    default_setting['Method'] = 'Flux' if st.sidebar.selectbox(
        "燃耗方式", ["定通量燃耗", "定功率燃耗"]) == "定通量燃耗" else 'Power'
    if default_setting['Method'] == 'Flux':
        default_setting['Flux'] = st.sidebar.number_input(
            "中子通量（每平方厘米每秒）", min_value=0.0, format='%e', value=default_setting['Flux'], step=1e13)
        st.title("同位素成分随时间变化")
    else:
        default_setting['Power'] = st.sidebar.number_input(
            '比功率（兆瓦每吨）', min_value=0.0, format='%e', value=float(default_setting['Power']), help="单位质量的燃料可以放出的功率", step=50.0)
        st.title("同位素成分随燃耗变化")

    for value in DATA.values():
        st.write(value)

    default_setting['t'] = st.number_input(
        "运行时间（天）", min_value=0, format='%d', value=default_setting['t'], step=30)
    default_setting['DT'] = st.slider(
        "燃耗步长（天）", min_value=1, max_value=30, format='%d', value=default_setting['DT'])
    PLOT_ACCURACY = st.number_input("内步长精度", min_value=1e-8, value=default_setting['Accuracy'], format="%e", step=1e-4)

    draw_names = []
    DRAW_NAMES = []
    draw_list = {}

    SUM_INITIAL = 0
    for name, nuclides in default_setting['Initial'].items():
        draw_list[name] = {}
        for nuclide, weight in nuclides.items():
            draw_list[name][nuclide] = {
                "Enable": False,
                "Ratio": 1
            }
            draw_names.append(nuclide)
            if DRAW_LIST[name][nuclide]['Enable']:
                DRAW_NAMES.append(nuclide)
            if nuclide in ['234U', '235U', '238U']:
                default_setting['Initial'][name][nuclide] = st.number_input(
                    label=nuclide + "含量（千克每吨）", min_value=0.0, max_value=1000.0, value=float(weight), format='%f', help=nuclide + "在单位质量燃料中的含量", step=1.0)
            SUM_INITIAL += default_setting['Initial'][name][nuclide]

    if SUM_INITIAL > 1000:
        st.warning("初始成分之和超出范围。请更正数值。")
    else:
        select_result = st.multiselect("需要展示的核素", draw_names, default=DRAW_NAMES)
        for name, nuclides in draw_list.items():
            for nuclide, plot_conf in nuclides.items():
                if nuclide in select_result:
                    plot_conf['Enable'] = True
                    plot_conf['Ratio'] = float(st.sidebar.number_input(
                        nuclide + "缩放比例", min_value=0.0, format='%e', value=float(DRAW_LIST[name][nuclide]['Ratio']), step=0.1))

        PLOT_MODE = st.selectbox("绘图方式", ["解析方法", "数值方法", "两种方法"])

        if PLOT_MODE in ["解析方法", "两种方法"]:
            start = time.time()
            analytic = Analytic(DATA_Cal, default_setting)
            analytic.main(default_setting['t'] * 24 * 60 * 60, PLOT_ACCURACY)
            stop = time.time()
            plot = Plot(analytic.numerical_result)
            plot.Draw(draw_list)
            plt.ion()
            st.pyplot(plot.FIG)
            st.info("解析方法计算用时%.3e秒。" % (stop-start))

        if PLOT_MODE in ["数值方法", "两种方法"]:
            start = time.time()
            numerical = Numerical(DATA_Cal, default_setting)
            numerical.main(default_setting['t'] * 24 * 60 * 60, PLOT_ACCURACY)
            stop = time.time()
            plot = Plot(numerical.numerical_result)
            plot.Draw(draw_list)
            st.pyplot(plot.FIG)
            st.info("数值方法计算用时%.3e秒。" % (stop-start))

if __name__ == '__main__':
    DATA = {}
    for CSV_File in DATASETS:
        with open(CSV_File, encoding='gbk') as csv_file:
            DATA[CSV_File[-8:-4]] = pd.read_csv(csv_file)
    with open(SETTINGS[-1], 'r') as f:
        default_setting = json.load(f)

    Setting_Mode = st.sidebar.selectbox("配置模式", ["默认模式", "专家模式"])

    Uploaded_Files = {}
    DATA_Cal = {}
    Is_Uploaded = True
    if Setting_Mode == "专家模式":
        for key, value in DATA.items():
            csv = value.to_csv(index=False)
            b64 = base64.b64encode(csv.encode(encoding='gbk')).decode(encoding='gbk')
            href = f'<a href="data:file/csv;base64,{b64}">下载{key}燃耗链配置文件</a> (右键点击链接并另存为csv格式)'
            st.markdown(href, unsafe_allow_html=True)
            Uploaded_Files[key] = st.file_uploader("选择您的配置文件", type="csv", key=key + '_upload')
        for key, value in Uploaded_Files.items():
            if value is not None:
                data_pd = pd.read_csv(value, encoding='gbk')
                DATA[key] = data_pd
                DATA_Cal[key] = data_pd.to_numpy()
            else:
                Is_Uploaded = False
        if Is_Uploaded:
            main(DATA_Cal, DATA)
    else:
        for key, value in DATA.items():
            DATA_Cal[key] = value.to_numpy()
        main(DATA_Cal, DATA)
